#

import time
import unittest
import yaml
from selenium import webdriver


class TestCookieLogin(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.driver = webdriver.Chrome()

    @classmethod
    def tearDownClass(cls):
        cls.driver.quit()

    def test_get_cookies(self):
        self.driver.get("http://work.weixin.qq.com/wework_admin/frame")
        time.sleep(20)
        cookies = self.driver.get_cookies()
        print(cookies)
        # 修复：正确写入文件，使用 with 确保文件关闭
        with open("cookie.yaml", "w") as f:
            yaml.safe_dump(cookies, f)  # 正确写法：yaml.safe_dump(数据, 文件对象)

    def test_add_cookies(self):
        # 先访问一次页面，确保域名一致才能添加 cookie
        self.driver.get("http://work.weixin.qq.com/wework_admin/frame")
        cookies = yaml.safe_load(open("cookie.yaml"))
        # # 修复：检查文件是否存在且内容有效
        try:
            with open("cookie.yaml", "r") as f:
                cookies = yaml.safe_load(f)
                if cookies is None:
                    self.fail("cookie.yaml 文件内容为空或格式错误")
        except FileNotFoundError:
            self.fail("cookie.yaml 文件不存在，请先运行 test_get_cookies")

        # 添加 cookie
        for c in cookies:
            self.driver.add_cookie(c)
        time.sleep(3)
        # 刷新页面验证登录状态
        self.driver.refresh()
        self.driver.get("http://work.weixin.qq.com/wework_admin/frame")